from .account import Account
from .pyechart_utils import PyechartUtils


class Strategy:
    def __init__(self, feed, algo_list=None):
        self.name = ""
        self.feed = feed
        self.algo_list = algo_list
        self.acc = Account()

        # 从feed得到df，并从index取出日期列表
        self.df = feed.get_all_df()
        self.dates = self.df.index.unique()

    def run(self):
        for index, date in enumerate(self.dates):
            self.step(index, date)

    def step(self, index, date):
        df_bar = self.df.loc[date]
        if type(df_bar) is pd.Series:
            df_bar = df_bar.to_frame().T

        df_bar.index = df_bar['code']
        df_bar.sort_index(ascending=True, inplace=True)
        self.acc.update_bar(date,df_bar)

        self.onbar(index, date, df_bar)

    def onbar2(self, index, date, df_bar):

        self.acc.update_bar(date, df_bar)
        weights = self.algo_processor({'index': index, 'bar': df_bar, 'date': date, 'acc': self.acc})
        if weights is not None:
            self.acc.adjust_weights(date, weights)

    def algo_processor(self, context):
        for algo in self.algo_list:
            if algo(context) is True:  # 如果algo返回True,直接不运行，本次不调仓
                return None
        return context['weights']

    def plot(self, benchmark=None, filename='strategy.html'):
        df = self.acc.get_results_df()

        df['策略'] = df['rate']
        df['策略_equity'] = (1 + df['rate']).cumprod()
        equites = df[['策略_equity']]

        rates_cols = ['策略']

        for bench in benchmark:
            df_bench = self.feed.get_df(bench)
            df["基准_" + bench] = df_bench['rate']
            rates_cols.append("基准_" + bench)
            equites['基准_' + bench] = df_bench['equity']

        from engine.performance import PerformanceUtils

        df_ratios, df_corr, df_years = PerformanceUtils().calc_rates(df[rates_cols])
        print(df_ratios)

        echarts = PyechartUtils()
        echarts.add_lines(equites, '回测曲线')
        echarts.add_table(df_ratios, '风险收益特征')
        echarts.render(filename)


from engine.algos import *


class StrategyBuyHold(Strategy):
    # params={period=1...N,可选；}
    def __init__(self, name, params={}):
        algo_period = None
        if 'period' in params.keys():
            algo_period = RunPeriod(period=params['period'])
        else:
            algo_period = RunOnce()
        algo_list = [
            algo_period,
            SelectAll(),
            WeightEqually()
        ]
        super(StrategyBuyHold, self).__init__(name, algo_list)


class StrategyRolling(Strategy):

    def __init__(self, name, params={}):
        algo_list = [
            SelectBySignal(signal_buy=params['signal_buy'], signal_sell=params['signal_sell']),
            SelectTopK(K=params['K'], col=params['sort_by']),
            WeightEqually()
        ]
        super(StrategyRolling, self).__init__(name, algo_list=algo_list)


class StrategyMachineLearning(Strategy):
    def __init__(self, name, params={}):
        algo_list = [
            # SelectBySignal(signal_buy=params['signal_buy'],signal_sell=params['signal_sell']),
            SelectTopK(K=params['K'], col=params['sort_by']),
            WeightEqually()
        ]
        super(StrategyMachineLearning, self).__init__(name, algo_list=algo_list)
        self._init_model(params)

    def _init_model(self, params):
        df = params['df']
        from engine.model.gbdt import LGBModel
        m = LGBModel()
        m.fit(df)
        results = m.predict()
        df['pred_score'] = results
